%% Reproducing figures for Linz et al. 2021
% performing calculations of various metrics
clear all
close all

load 'WACCM_pvars1.mat'
load 'WACCM_pvars2.mat'
load 'WACCM_pvars3.mat'


g=9.8;
rovercp=0.286;
R=287; %J/kg/K
ps=10^5;%surface pressure
r=6378100; %radius of earth in m
resx=2.5;
resy=1.8947;

load('_interpp_log.mat')
load('_interpQ_log.mat')
load('_interp_isend_log.mat')

dwnmsk2=zeros(size(squeeze(mean(Q_interp1,4))));
dwnmsk2(squeeze(mean(Q_interp1,4))<0)=1;
upmsk2=zeros(size(squeeze(mean(Q_interp1,4))));
upmsk2(squeeze(mean(Q_interp1,4))>=0)=1;
dwnmsk2=repmat(dwnmsk2,[1 1 1 419]);
upmsk2=repmat(upmsk2, [1 1 1 419]);
dwnmskNH=dwnmsk2;
dwnmskSH=dwnmsk2;
dwnmskNH(:,1:47,:,:)=0;
dwnmskSH(:,48:end,:,:)=0;

% msk2 = time average upwelling/downwelling region
% dwnmskNH, SH can define the downwelling for each hemisphere, but need to
% figure out equator for upwelling.

%%%% upwelling mass, downwelling mass

Muup2=squeeze(sum(sum(weightedsigmatheta.*upmsk2,1),2));

Mudwn2=squeeze(sum(sum(weightedsigmatheta.*dwnmsk2,1),2));


% hems down
MudwnNH=squeeze(sum(sum(weightedsigmatheta.*dwnmskNH,1),2));
MudwnSH=squeeze(sum(sum(weightedsigmatheta.*dwnmskSH,1),2));

% hems down time avg.
MudwnNHavg=squeeze(mean(MudwnNH(:,1:408),2));
MudwnSHavg=squeeze(mean(MudwnSH(:,1:408),2));
% 


% where is the dynamical equator?
miniSH=zeros(54,30);
miniNH=zeros(54,30);
    for nlat=20:73
        avgupSH=upmsk2;
        avgupSH(:,nlat:end,:)=0;
        MuSHup_=squeeze(sum(sum(mean(weightedsigmatheta.*avgupSH,4),1),2));
        miniSH(nlat-19,:)=MuSHup_+MudwnSHavg;
                avgupNH=upmsk2;
        avgupNH(:,1:nlat,:)=0;
        MuNHup_=squeeze(sum(sum(mean(weightedsigmatheta.*avgupNH,4),1),2));
        miniNH(nlat-19,:)=MuNHup_+MudwnNHavg;
    end
    latinds=zeros(30,1);
    latindn=zeros(30,1);
    avgupSH=upmsk2;
    avgupNH=upmsk2;
    
        avgupSHn=upmsk2;
    avgupNHn=upmsk2;
for llev=1:30
    [M,I]=min(abs(miniSH(:,llev)));
    latinds(llev)=I;
    avgupSH(:,19+I:end,llev,:)=0;
    avgupNH(:,1:19+I-1,llev,:)=0;
    [M,I]=min(abs(miniNH(:,llev)));
    latindn(llev)=I;
    avgupSHn(:,19+I+1:end,llev,:)=0;
    avgupNHn(:,1:19+I,llev,:)=0;
end
avgupSHf=avgupSH;
avgupSHf(avgupSH-avgupSHn==1)=0.5;
avgupNHf=avgupNHn;
avgupNHf(avgupNHn-avgupNH==1)=0.5;


Muupsh=squeeze(sum(sum(weightedsigmatheta.*avgupSHf,1),2));
Muupnh=squeeze(sum(sum(weightedsigmatheta.*avgupNHf,1),2));


%%%% Total overturning
Mutotal2 = 1/2 * (Muup2 - Mudwn2);

load('WACCM_free_age.mat')

 age=repmat(age_interp, [1 1 1 144]);
    age=permute(age, [4 1 2 3]);

    ageup2=squeeze(sum(sum(weightedsigmatheta.*age.*upmsk2,1),2))./Muup2;
    agedwn2=squeeze(sum(sum(weightedsigmatheta.*age.*dwnmsk2,1),2))./Mudwn2;
    
    ageupNH=mean(squeeze(sum(sum(weightedsigmatheta.*age.*avgupNHf,1),2)),2)./squeeze(mean(Muupnh,2));
    ageupSH=mean(squeeze(sum(sum(weightedsigmatheta.*age.*avgupSHf,1),2)),2)./squeeze(mean(Muupsh,2));
   
    agedwnNH=mean(squeeze(sum(sum(weightedsigmatheta.*age.*dwnmskNH,1),2)),2)./MudwnNHavg;
    agedwnSH=mean(squeeze(sum(sum(weightedsigmatheta.*age.*dwnmskSH,1),2)),2)./MudwnSHavg;

    totalmassNH=squeeze(nansum(nansum(weightedp.*(avgupNHf+dwnmskNH)*100,1),2))/9.8;
    totalmassSH=squeeze(nansum(nansum(weightedp.*(avgupSHf+dwnmskSH)*100,1),2))/9.8;
       
    MutotNH=(Muupnh-MudwnNH)/2;
    MutotSH=(Muupsh-MudwnSH)/2;
    Mutotal2=1/2*(Muup2-Mudwn2);
    dgamma2=agedwn2-ageup2;
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        % vertical gradient theory calculations
    sigmaup=squeeze(sum(sum(weightedsigma.*upmsk2,1),2));
    sigmadwn=squeeze(sum(sum(weightedsigma.*dwnmsk2,1),2));
    sigmaupNH=mean(squeeze(sum(sum(weightedsigma.*avgupNHf,1),2)),2);
    sigmadwnNH=mean(squeeze(sum(sum(weightedsigma.*dwnmskNH,1),2)),2);
    sigmaupSH=mean(squeeze(sum(sum(weightedsigma.*avgupSHf,1),2)),2);
    sigmadwnSH=mean(squeeze(sum(sum(weightedsigma.*dwnmskSH,1),2)),2);
    intsigthdotgamma=squeeze(sum(sum(weightedsigmatheta.*age.*upmsk2,1),2));
  

    dgammaupdtheta=zeros(size(ageup2));
    dgammadwndtheta=zeros(size(ageup2));
    dgammaupdthetaNH=zeros(size(ageupNH));
    dgammadwndthetaNH=zeros(size(ageupNH));
    dgammaupdthetaSH=zeros(size(ageupNH));
    dgammadwndthetaSH=zeros(size(ageupNH));
    ddth_upflux=zeros(size(ageup2));
    for ll=2:29
        dgammaupdtheta(ll,:) = (ageup2(ll+1,:)-ageup2(ll-1,:))*365*24*60*60/(Thlevels(ll+1)-Thlevels(ll-1));
        dgammadwndtheta(ll,:) = (agedwn2(ll+1,:)-agedwn2(ll-1,:))*365*24*60*60/(Thlevels(ll+1)-Thlevels(ll-1));
        dgammaupdthetaNH(ll) = (ageupNH(ll+1)-ageupNH(ll-1))*365*24*60*60/(Thlevels(ll+1)-Thlevels(ll-1));
        dgammadwndthetaNH(ll) = (agedwnNH(ll+1)-agedwnNH(ll-1))*365*24*60*60/(Thlevels(ll+1)-Thlevels(ll-1));
        dgammaupdthetaSH(ll) = (ageupSH(ll+1)-ageupSH(ll-1))*365*24*60*60/(Thlevels(ll+1)-Thlevels(ll-1));
        dgammadwndthetaSH(ll) = (agedwnSH(ll+1)-agedwnSH(ll-1))*365*24*60*60/(Thlevels(ll+1)-Thlevels(ll-1));
        ddth_upflux(ll,:) = (ageup2(ll+1,:).*Mutotal2(ll+1,:)-ageup2(ll-1,:).*Mutotal2(ll-1,:))*365*24*60*60/(Thlevels(ll+1)-Thlevels(ll-1));
    end
    dgammaupdtheta(1,:)=(ageup2(2,:)-ageup2(1,:))*365*24*60*60/2/(Thlevels(2)-Thlevels(1));
    dgammaupdtheta(30,:)=(ageup2(30,:)-ageup2(29,:))*365*24*60*60/2/(Thlevels(30)-Thlevels(29));
    
    dgammadwndtheta(1,:)=(agedwn2(2,:)-agedwn2(1,:))*365*24*60*60/2/(Thlevels(2)-Thlevels(1));
    dgammadwndtheta(30,:)=(agedwn2(30,:)-agedwn2(29,:))*365*24*60*60/2/(Thlevels(30)-Thlevels(29));
    
    dgammaupdthetaNH(1)=(ageupNH(2)-ageupNH(1))*365*24*60*60/2/(Thlevels(2)-Thlevels(1));
    dgammaupdthetaNH(30)=(ageupNH(30)-ageupNH(29))*365*24*60*60/2/(Thlevels(30)-Thlevels(29));
    
    dgammadwndthetaNH(1)=(agedwnNH(2)-agedwnNH(1))*365*24*60*60/2/(Thlevels(2)-Thlevels(1));
    dgammadwndthetaNH(30)=(agedwnNH(30)-agedwnNH(29))*365*24*60*60/2/(Thlevels(30)-Thlevels(29));
    
    dgammaupdthetaSH(1)=(ageupSH(2)-ageupSH(1))*365*24*60*60/2/(Thlevels(2)-Thlevels(1));
    dgammaupdthetaSH(30)=(ageupSH(30)-ageupSH(29))*365*24*60*60/2/(Thlevels(30)-Thlevels(29));
    
    dgammadwndthetaSH(1)=(agedwnSH(2)-agedwnSH(1))*365*24*60*60/2/(Thlevels(2)-Thlevels(1));
    dgammadwndthetaSH(30)=(agedwnSH(30)-agedwnSH(29))*365*24*60*60/2/(Thlevels(30)-Thlevels(29));
    
    ddth_upflux(1,:)=(ageup2(2,:).*Mutotal2(2,:)-ageup2(1,:).*Mutotal2(1,:))*365*24*60*60/2/(Thlevels(2)-Thlevels(1));
    ddth_upflux(30,:)=(ageup2(30,:).*Mutotal2(30,:)-ageup2(29,:).*Mutotal2(29,:))*365*24*60*60/2/(Thlevels(30)-Thlevels(29));
    
    muoversigdgdth=dgammaupdtheta.*Mutotal2./sigmaup;
    
    muoutoversig=mean(Mutotal2,2).*mean(dgammadwndtheta,2)./mean(sigmadwn,2)./mean(agedwn2-ageup2,2)+1./mean(agedwn2-ageup2,2);
    muinoversig=mean(Mutotal2,2).*mean(dgammaupdtheta,2)./mean(sigmaup,2)./mean(agedwn2-ageup2,2)-1./mean(agedwn2-ageup2,2);
    
    muoutoversigNH=squeeze(mean(MutotNH,2)).*dgammadwndthetaNH./sigmadwnNH./(agedwnNH-ageupNH)+1./(agedwnNH-ageupNH);
    muinoversigNH=squeeze(mean(MutotNH,2)).*dgammaupdthetaNH./sigmaupNH./(agedwnNH-ageupNH)-1./(agedwnNH-ageupNH);
    
    muoutoversigSH=squeeze(mean(MutotSH,2)).*(dgammadwndthetaSH)./(sigmadwnSH)./(agedwnSH-ageupSH)+1./(agedwnSH-ageupSH);
    muinoversigSH=squeeze(mean(MutotSH,2)).*(dgammaupdthetaSH)./(sigmaupSH)./(agedwnSH-ageupSH)-1./(agedwnSH-ageupSH);
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    % now load in the idealized model data set
    
    M = zeros(17,25);
MNH = zeros(17, 25);
MSH = zeros(17, 25);
gammau = zeros(17, 25);
gammad = zeros(size(M));
gammadNH = zeros(size(M));
gammadSH = zeros(size(M));
gammauNH = zeros(size(M));
gammauSH = zeros(size(M));
muupNH = zeros(size(M));
mudwnNH = zeros(size(M));
muupSH = zeros(size(M));
mudwnSH = zeros(size(M));
dgudth_ideal = zeros(size(M));

ideal_levels = ncread('thdot_09490c.nc','isen');

    filenum = 9490; 
for nyr = 1:2
    
M(:,nyr) = ncread(strcat('thdot_0', num2str(filenum), 'c.nc'),'M');
MNH(:,nyr) = ncread(strcat('thdot_0', num2str(filenum), 'c.nc'),'MNH');
MSH(:,nyr) = ncread(strcat('thdot_0', num2str(filenum), 'c.nc'),'MSH');
gammau(:,nyr) = ncread(strcat('thdot_0', num2str(filenum), 'c.nc'),'gammau');
gammad(:,nyr) = ncread(strcat('thdot_0', num2str(filenum), 'c.nc'),'gammad');
gammauNH(:,nyr) = ncread(strcat('thdot_0', num2str(filenum), 'c.nc'),'gammauNH');
gammauSH(:,nyr) = ncread(strcat('thdot_0', num2str(filenum), 'c.nc'),'gammauSH');
gammadNH(:,nyr) = ncread(strcat('thdot_0', num2str(filenum), 'c.nc'),'gammadNH');
gammadSH(:,nyr) = ncread(strcat('thdot_0', num2str(filenum), 'c.nc'),'gammadSH');
muupNH(:,nyr) = ncread(strcat('thdot_0', num2str(filenum), 'c.nc'),'muupNH');
muupSH(:,nyr) = ncread(strcat('thdot_0', num2str(filenum), 'c.nc'),'muupSH');
mudwnNH(:,nyr) = ncread(strcat('thdot_0', num2str(filenum), 'c.nc'),'mudwnNH');
mudwnSH(:,nyr) = ncread(strcat('thdot_0', num2str(filenum), 'c.nc'),'mudwnSH');
dgudth_ideal(:,nyr) = ncread(strcat('thdot_0', num2str(filenum), 'c.nc'),'dgupdth');

    filenum = filenum + 365
end
filenum = 10220; 
for nyr = 3:25
    
M(:,nyr) = ncread(strcat('thdot_', num2str(filenum), 'c.nc'),'M');
MNH(:,nyr) = ncread(strcat('thdot_', num2str(filenum), 'c.nc'),'MNH');
MSH(:,nyr) = ncread(strcat('thdot_', num2str(filenum), 'c.nc'),'MSH');
gammau(:,nyr) = ncread(strcat('thdot_', num2str(filenum), 'c.nc'),'gammau');
gammad(:,nyr) = ncread(strcat('thdot_', num2str(filenum), 'c.nc'),'gammad');
gammauNH(:,nyr) = ncread(strcat('thdot_', num2str(filenum), 'c.nc'),'gammauNH');
gammauSH(:,nyr) = ncread(strcat('thdot_', num2str(filenum), 'c.nc'),'gammauSH');
gammadNH(:,nyr) = ncread(strcat('thdot_', num2str(filenum), 'c.nc'),'gammadNH');
gammadSH(:,nyr) = ncread(strcat('thdot_', num2str(filenum), 'c.nc'),'gammadSH');
muupNH(:,nyr) = ncread(strcat('thdot_', num2str(filenum), 'c.nc'),'muupNH');
muupSH(:,nyr) = ncread(strcat('thdot_', num2str(filenum), 'c.nc'),'muupSH');
mudwnNH(:,nyr) = ncread(strcat('thdot_', num2str(filenum), 'c.nc'),'mudwnNH');
mudwnSH(:,nyr) = ncread(strcat('thdot_', num2str(filenum), 'c.nc'),'mudwnSH');
dgudth_ideal(:,nyr) = ncread(strcat('thdot_', num2str(filenum), 'c.nc'),'dgupdth');
    filenum = filenum + 365
end

dgamma = gammad-gammau;
dgammaSH = gammadSH - gammauSH;
dgammaNH = gammadNH - gammauNH;

mutot = 1/2 * (muupSH-mudwnSH) + 1/2 * (muupNH - mudwnNH);

Movermu = M./mutot;
MovermuNH = 2*MNH./(muupNH - mudwnNH);
MovermuSH = 2*MSH./(muupSH - mudwnSH);


%% Figure 1: Diagram
%% Figure 2: AoA and streamfunctions for idealized model, WACCM, N2O
% Fig 2c: N2O_age

load('N2O_age_andrews.mat')
figure
contourf(lat,Thlevels,squeeze(mean(gammalag_andrews_th(:,:,12:end),3)'),0:.5:7,'linewidth',0)
colormap(parula(14))
set(gca,'fontsize',20,'Clim', [0 7],'yscale','log')
set(gca,'Ylim',[400 1000])
xlabel('latitude (^{\circ}N)')
ylabel('\Theta (K)')
colorbar
saveas(gcf,'Linzetal2021_AoA_Fig2c','pdf')

% Fig 2a: idealized model
load('aman_tac_age_streamfunction_levels.mat')
load('aman_tac_coords.mat')
y=logspace(0,3,10);
v=[-1000 -100 -10 -1 0  1 10 100 1000];
thlevels=isen;
contours_new=[-y 0 y];
latweight1=linspace(0,1,92);
latweight2=linspace(1,0,92);
latweight1=repmat(latweight1', [1 17]);
latweight2=repmat(latweight2',[1 17]);
figure
contourf(lat,isen,meanage,0:.5:7,'linewidth',0)
colormap(parula(14))
set(gca,'fontsize',20,'Clim', [0 7],'yscale','log')
set(gca,'Ylim',[400 1000])
xlabel('latitude (^{\circ}N)')
ylabel('\Theta (K)')
colorbar
hold on
[C,h]=contour(lat, isen, (-streamfunctionN'.*latweight1'+streamfunctionS'.*latweight2')*r,contours_new,'k');
clabel(C,h,v,'fontsize',16)
saveas(gcf,'Linzetal2021_AoA_Fig2a','pdf')


% Fig 2b: WACCM
load 'WACCM_free_streamfunction.mat'
load 'WACCM_coords.mat'
r=6378100;% radius of earth in m

y=logspace(0,3,10);
v=[-1000 -100 -10 -1 0  1 10 100 1000];

contours_new=[-y 0 y];
latweight1=linspace(0,1,96);
latweight2=linspace(1,0,96);
latweight1=repmat(latweight1', [1 30]);
latweight2=repmat(latweight2',[1 30]);
figure
contourf(lat,Thlevels(1:26),meanage(1:26,:),0:.5:7,'linewidth',0)
colormap(parula(14))
set(gca,'fontsize',20,'Clim', [0 7],'yscale','log')
set(gca,'Ylim',[400 1000])
xlabel('latitude (^{\circ}N)')
ylabel('\Theta (K)')
colorbar
hold on
[C,h]=contour(lat, Thlevels(1:26), (-streamfunctionN(:,1:26)'.*latweight1(:,1:26)'+streamfunctionS(:,1:26)'.*latweight2(:,1:26)')*r,contours_new,'k');
clabel(C,h,v,'fontsize',16)
saveas(gcf,'Linzetal2021_AoA_Fig2b','pdf')

%% Figure 3: Scatter plots idealized model and WACCM-time averaging

% Fig 3b: 1 yr vs 5 yr averaging, WACCM 
%*****%load('WACCM_free_dgammaetc.mat','dgamma2','totalmass','Mutotal2','Thlevels'
% 1 year averages WACCM -- note that these use the time averaged upwelling
% and downwelling regions


    dgammayravg=zeros(30,34);
    Movermuyravg=zeros(30,34);
    
    for yr=1:34
    dgammayravg(:,yr)=mean(dgamma2(:,(yr-1)*12+12:yr*12+11),2);
    Movermuyravg(:,yr)=mean((WACCMtotalmass(:,(yr-1)*12+12:yr*12+11)./Mutotal2(:,(yr-1)*12+12:yr*12+11)),2);
    end

% 5 year averages WACCM -- note that these use the time averaged upwelling
% and downwelling regions
    
    dgamma5yravg=zeros(30,6);
    Movermu5yravg=zeros(30,6);
    
    for yr=1:6
    dgamma5yravg(:,yr)=mean(dgamma2(:,(yr-1)*5*12+12:yr*5*12+11),2);
    Movermu5yravg(:,yr)=mean((WACCMtotalmass(:,(yr-1)*12*5+12:yr*5*12+11)./Mutotal2(:,(yr-1)*12*5+12:yr*12*5+11)),2);
    end


    colors1=parula(30);
    figure
    hold on
    countn=0;
    for ll=2:2:30
        countn=countn+2;
       scatter(dgammayravg(ll,:),Movermuyravg(ll,:)/365/24/60/60,80,'MarkerEdgeColor',colors1(countn,:))
       hold on
    end
    plot(0.4:.1:2,0.4:.1:2,'k')
    set(gca,'fontsize',18)
    xlabel('\Delta\Gamma')

    ylabel('mass above \theta/ mass flux across \theta (years)')
    legend(num2str(Thlevels(2)),num2str(Thlevels(4)),num2str(Thlevels(6)),num2str(Thlevels(8)),num2str(Thlevels(10)),num2str(Thlevels(12)),...
        num2str(Thlevels(14)),num2str(Thlevels(16)),num2str(Thlevels(18)),num2str(Thlevels(20)),num2str(Thlevels(22)),num2str(Thlevels(24)),...
        num2str(Thlevels(26)),num2str(Thlevels(28)),num2str(Thlevels(30)),'location','northwest');
     countn=0;
    for ll=2:2:30
        countn=countn+2;
       scatter(dgamma5yravg(ll,:),Movermu5yravg(ll,:)/365/24/60/60,150,'*','MarkerEdgeColor',colors1(countn,:),'linewidth',3)
       hold on
    end
    title('5 year averaged idealized, colored by \Theta')
    set(gca,'xlim', [0.5 2], 'ylim', [0.5 2])
set(gcf, 'PaperUnits', 'inches'); set(gcf, 'PaperSize', [6 6]);

set(gcf, 'PaperPosition', [0 0 6 6]);
saveas(gcf,'Linzetal2021_AoA_Fig3b','pdf')
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


% Fig 3a: 1 yr vs 5 yr averaging, idealized model
load('aman_tac_coords.mat')
% 5 year averages
Movermuavg = zeros(17,5);

dgammaavg = zeros(17,5);
thlevels= isen;

for nn = 1:5
    Movermuavg(:,nn) = mean(Movermu(:,(nn-1)*5+1:1:5*nn),2);

    dgammaavg(:,nn) = mean(dgamma(:,(nn-1)*5+1:1:5*nn),2);

end
 addpath cbrewer
    %colors1=cbrewer('seq','GnBu',35);
    colors1=parula(12);
    figure
    hold on
    countn=0;
    for ll=[2 4 6 9 10:17]
        countn=countn+1;
       scatter(dgamma(ll,:),Movermu(ll,:),80,'MarkerEdgeColor',colors1(countn,:))
       hold on
    end
    plot(1.2:.1:2.3,1.2:.1:2.3,'k')
    set(gca,'fontsize',18)
    xlabel('\Delta\Gamma')
    %colorbar('Ticks',[0,.25, .5, .75, 1],...
    %    'TickLabels',{'400','525', '700', '950', '1200'})
    
    ylabel('mass above \theta/ mass flux across \theta (years)')
    legend(num2str(thlevels(2)),num2str(thlevels(4)),num2str(thlevels(6)),...
        num2str(thlevels(9)),num2str(thlevels(10)),num2str(thlevels(11)),num2str(thlevels(12)),...
        num2str(thlevels(13)),num2str(thlevels(14)),num2str(thlevels(15)),num2str(thlevels(16)),...
        num2str(thlevels(17)),'location','northwest');
     countn=0;
    for ll=[2 4 6 9 10:17]
        countn=countn+1;
       scatter(dgammaavg(ll,:),Movermuavg(ll,:),150,'*','MarkerEdgeColor',colors1(countn,:),'linewidth',3)
       hold on
    end
    title('5 year averaged idealized, colored by \Theta')
    set(gca,'xlim', [1.2 2.3], 'ylim', [1.2 2.3])
set(gcf, 'PaperUnits', 'inches'); set(gcf, 'PaperSize', [6 6]);

set(gcf, 'PaperPosition', [0 0 6 6]);
saveas(gcf,'Linzetal2021_AoA_Fig3b','pdf')


%% Figure 4: Ideal and WACCM hemispheric separation 
% age difference vs M over Mu for both models separated by hemisphere (long averages)
colors2 = parula(30);
figure
for countn = 2:2:30
h1a=scatter(mean(agedwn2(countn,:)-ageup2(countn,:),2),squeeze(mean((WACCMtotalmass(countn,:)./Mutotal2(countn,:)),2))/365/24/60/60,150,'*','MarkerEdgeColor',colors2(countn,:),'linewidth',2)
hold on
end
for countn=2:2:30
h2a=scatter(mean(agedwnNH(countn,:)-ageupNH(countn,:),2),squeeze(mean((totalmassNH(countn,:)./MutotNH(countn,:)),2))/365/24/60/60,100,'s','MarkerEdgeColor',colors2(countn,:),'linewidth',2)
hold on
h3a=scatter(mean(agedwnSH(countn,:)-ageupSH(countn,:),2),squeeze(mean((totalmassSH(countn,:)./MutotSH(countn,:)),2))/365/24/60/60,100,'d','MarkerEdgeColor',colors2(countn,:),'linewidth',2)

end
plot(0.5:.1:2,0.5:.1:2,'k','linestyle','--')
set(gca,'fontsize',18,'xlim',[.5 2],'ylim',[.5 2])
 legend(num2str(Thlevels(2)),num2str(Thlevels(4)),num2str(Thlevels(6)),num2str(Thlevels(8)),num2str(Thlevels(10)),num2str(Thlevels(12)),...
        num2str(Thlevels(14)),num2str(Thlevels(16)),num2str(Thlevels(18)),num2str(Thlevels(20)),num2str(Thlevels(22)),num2str(Thlevels(24)),...
        num2str(Thlevels(26)),num2str(Thlevels(28)),num2str(Thlevels(30)),'location','southeast')
    ylabel('mass above \theta/ mass flux across \theta (years)')
    xlabel('\Delta \Gamma(\theta) (years)')
             set(gcf, 'PaperUnits', 'inches'); set(gcf, 'PaperSize', [6 6]);

set(gcf, 'PaperPosition', [0 0 6 6]);
saveas(gcf,'Linzetal2021_AoA_Fig4b','pdf')


% Same for idealized model 
figure
colors3 = parula(17)
for countn = 1:17
h1b=scatter(mean(dgamma(countn,:),2),squeeze(mean(Movermu(countn,:),2)),150,'*','MarkerEdgeColor',colors3(countn,:),'linewidth',2)
hold on
end
for countn = 1:17
h2b=scatter(mean(dgammaNH(countn,:),2),squeeze(mean(MovermuNH(countn,:),2)),100,'s','MarkerEdgeColor',colors3(countn,:),'linewidth',2)
h3b=scatter(mean(dgammaSH(countn,:),2),squeeze(mean(MovermuSH(countn,:),2)),100,'d','MarkerEdgeColor',colors3(countn,:),'linewidth',2)
end
plot(1.2:.1:2.5,1.2:.1:2.5,'k','linestyle','--')
set(gca,'fontsize',18,'xlim',[1.2 2.5],'ylim',[1.2 2.5])

    ylabel('mass above \theta/ mass flux across \theta (years)')
    xlabel('\Delta \Gamma(\theta) (years)')
             set(gcf, 'PaperUnits', 'inches'); set(gcf, 'PaperSize', [6 6]);
 legend(num2str(ideal_levels(1)),num2str(ideal_levels(2)),num2str(ideal_levels(3)),num2str(ideal_levels(4)),num2str(ideal_levels(5)),num2str(ideal_levels(6)),...
        num2str(ideal_levels(7)),num2str(ideal_levels(8)),num2str(ideal_levels(9)),num2str(ideal_levels(10)),num2str(ideal_levels(11)),num2str(ideal_levels(12)),...
        num2str(ideal_levels(13)),num2str(ideal_levels(14)),num2str(ideal_levels(15)),num2str(ideal_levels(16)),num2str(ideal_levels(17)),'location','southeast')
set(gcf, 'PaperPosition', [0 0 6 6]);




saveas(gcf,'Linzetal2021_AoA_Fig4a','pdf')
%% Figure 5: Mixing rates into and out of the tropics
% Muin and Muout (each over sigma)

dgddth_ideal = zeros(size(M));
dgudth_idealNH = zeros(size(M));
dgudth_idealSH = zeros(size(M));
dgddth_idealNH = zeros(size(M));
dgddth_idealSH = zeros(size(M));
sigu = zeros(size(M));
sigd = zeros(size(M));
siguNH = zeros(size(M));
siguSH = zeros(size(M));
sigdNH = zeros(size(M));
sigdSH = zeros(size(M));


ideal_levels = ncread('thdot_09490c.nc','isen');

    filenum = 9490; 
for nyr = 1:2
    
dgddth_ideal(:,nyr) = ncread(strcat('thdot_0', num2str(filenum), 'c.nc'),'dgdwndth');
dgudth_idealNH(:,nyr) = ncread(strcat('thdot_0', num2str(filenum), 'c.nc'),'dgupdthNH');
dgudth_idealSH(:,nyr) = ncread(strcat('thdot_0', num2str(filenum), 'c.nc'),'dgupdthSH');
dgddth_idealNH(:,nyr) = ncread(strcat('thdot_0', num2str(filenum), 'c.nc'),'dgdwndthNH');
dgddth_idealSH(:,nyr) = ncread(strcat('thdot_0', num2str(filenum), 'c.nc'),'dgdwndthSH');
siguNH(:,nyr) = ncread(strcat('thdot_0', num2str(filenum), 'c.nc'),'sigmauNH');
siguSH(:,nyr) = ncread(strcat('thdot_0', num2str(filenum), 'c.nc'),'sigmauSH');
sigdNH(:,nyr) = ncread(strcat('thdot_0', num2str(filenum), 'c.nc'),'sigmadNH');
sigdSH(:,nyr) = ncread(strcat('thdot_0', num2str(filenum), 'c.nc'),'sigmadSH');
sigu(:,nyr) = ncread(strcat('thdot_0', num2str(filenum), 'c.nc'),'sigmau');
sigd(:,nyr) = ncread(strcat('thdot_0', num2str(filenum), 'c.nc'),'sigmad');

    filenum = filenum + 365
end
filenum = 10220; 
for nyr = 3:25
    
dgddth_ideal(:,nyr) = ncread(strcat('thdot_', num2str(filenum), 'c.nc'),'dgdwndth');
dgudth_idealNH(:,nyr) = ncread(strcat('thdot_', num2str(filenum), 'c.nc'),'dgupdthNH');
dgudth_idealSH(:,nyr) = ncread(strcat('thdot_', num2str(filenum), 'c.nc'),'dgupdthSH');
dgddth_idealNH(:,nyr) = ncread(strcat('thdot_', num2str(filenum), 'c.nc'),'dgdwndthNH');
dgddth_idealSH(:,nyr) = ncread(strcat('thdot_', num2str(filenum), 'c.nc'),'dgdwndthSH');
siguNH(:,nyr) = ncread(strcat('thdot_', num2str(filenum), 'c.nc'),'sigmauNH');
siguSH(:,nyr) = ncread(strcat('thdot_', num2str(filenum), 'c.nc'),'sigmauSH');
sigdNH(:,nyr) = ncread(strcat('thdot_', num2str(filenum), 'c.nc'),'sigmadNH');
sigdSH(:,nyr) = ncread(strcat('thdot_', num2str(filenum), 'c.nc'),'sigmadSH');
sigu(:,nyr) = ncread(strcat('thdot_', num2str(filenum), 'c.nc'),'sigmau');
sigd(:,nyr) = ncread(strcat('thdot_', num2str(filenum), 'c.nc'),'sigmad');
    filenum = filenum + 365
end

mutotNH = 1/2 * (muupNH - mudwnNH);
mutotSH = 1/2 * (muupSH-mudwnSH);

muinoversigmau = mutot./sigu./dgamma.*dgudth_ideal-1./dgamma;
muoutoversigmad = mutot./sigd./dgamma.*dgddth_ideal+1./dgamma;

muinoversigmauNH = mutotNH./siguNH./dgammaNH.*dgudth_idealNH-1./dgammaNH;
muoutoversigmadNH = mutotNH./sigdNH./dgammaNH.*dgddth_idealNH+1./dgammaNH;

muinoversigmauSH = mutotSH./siguSH./dgammaSH.*dgudth_idealSH-1./dgammaSH;
muoutoversigmadSH = mutotSH./sigdSH./dgammaSH.*dgddth_idealSH+1./dgammaSH;


figure
h6=plot(mean(muoutoversig(:,:),2),Thlevels(:),'color',[0.6 0.1 0.8], 'linewidth',3);
hold on
h7=plot(mean(muoutoversigNH(:,:),2),Thlevels(:), 'color',[0.6 0.1 0.8], 'linewidth',2,'linestyle','-.');
h8=plot(mean(muoutoversigSH(:,:),2),Thlevels(:), 'color',[0.6 0.1 0.8], 'linewidth',2,'linestyle',':');

h6a=plot(mean(muoutoversigmad(:,:),2),ideal_levels(:),'color',[0.03 0.58 0.13], 'linewidth',3);
hold on
h7a=plot(mean(muoutoversigmadNH(:,:),2),ideal_levels(:), 'color',[0.03 0.58 0.13], 'linewidth',2,'linestyle','-.');
h8a=plot(mean(muoutoversigmadSH(:,:),2),ideal_levels(:), 'color',[0.03 0.58 0.13], 'linewidth',2,'linestyle',':');

set(gca,'fontsize',18)

    ylabel('\theta (K)')
    xlabel('mixing rate into extratropics (1/yr)')
    [hleg,icons, thingy, stuffs]=legend([h6 h7 h8 h6a h7a h8a],{'WACCM global', 'WACCM NH','WACCM SH', 'ideal global','ideal NH','ideal SH'},'location','northeast','fontsize',18);
      
              set(gcf, 'PaperUnits', 'inches'); set(gcf, 'PaperSize', [5.5 6.5]);

set(gcf, 'PaperPosition', [0 0 5.5 6.5]);
set(gca,'yscale','log','ylim',[380 1200],'ytick',[400, 500, 600, 800, 1000, 1200])
saveas(gcf,'Linzetal2021_AoA_Fig5b','pdf')


figure
h6=plot(mean(muinoversig(:,:),2),Thlevels(:),'color',[0.6 0.1 0.8], 'linewidth',3);
hold on
h7=plot(mean(muinoversigNH(:,:),2),Thlevels(:), 'color',[0.6 0.1 0.8], 'linewidth',2,'linestyle','-.');
h8=plot(mean(muinoversigSH(:,:),2),Thlevels(:), 'color',[0.6 0.1 0.8], 'linewidth',2,'linestyle',':');

h6a=plot(mean(muinoversigmau(:,:),2),ideal_levels(:),'color',[0.03 0.58 0.13], 'linewidth',3);
hold on
h7a=plot(mean(muinoversigmauNH(:,:),2),ideal_levels(:), 'color',[0.03 0.58 0.13], 'linewidth',2,'linestyle','-.');
h8a=plot(mean(muinoversigmauSH(:,:),2),ideal_levels(:), 'color',[0.03 0.58 0.13], 'linewidth',2,'linestyle',':');
set(gca,'fontsize',18)

    ylabel('\theta (K)')
    xlabel('mixing rate into tropics (1/yr)')
  [hleg,icons, thingy, stuffs]=legend([h6 h7 h8 h6a h7a h8a],{'WACCM global', 'WACCM NH','WACCM SH', 'ideal global','ideal NH','ideal SH'},'location','northeast','fontsize',18);
       
              set(gcf, 'PaperUnits', 'inches'); set(gcf, 'PaperSize', [5.5 6.5]);

set(gcf, 'PaperPosition', [0 0 5.5 6.5]);
set(gca,'xlim',[-0.25 2])
set(gca,'yscale','log','ylim',[380 1200],'ytick',[400, 500, 600, 800, 1000, 1200])
saveas(gcf,'Linzetal2021_AoA_Fig5a','pdf')
%% Figure 6: Vertical gradient of upwelling age

%comparison of dgammaup/dtheta  both models and N2O
load('N2O_andrews_overturning.mat')
dgammaudth = zeros(size(tropage));
for ll=2:29
    dgammaudth(ll,:) = (tropage(ll+1,:)-tropage(ll-1,:))/(Thlevels(ll+1)-Thlevels(ll-1));
end
dgammaudthm = squeeze(nanmean(dgammaudth,2));
figure
h1=plot(mean(dgammaupdtheta(1:end,:),2)/365/24/60/60,Thlevels(1:end), 'color',[0.6 0.1 0.8], 'linewidth',3);
hold on
h2=plot(dgammaudthm(1:end),Thlevels(1:end),'color','k','linewidth',3);
h3 = plot(squeeze(mean(dgudth_ideal,2)),ideal_levels,'color',[0.03 0.58 0.13],'linewidth',3);
set(gca,'fontsize',18)
    ylabel('\theta (K)')
    xlabel('\partial \Gamma_u/\partial \theta (yr/K)')
   [hleg,icons, thingy, stuffs]=legend([h1 h2 h3],{'WACCM','N2O age', 'ideal'},'location','northeast','fontsize',18);
      
              set(gcf, 'PaperUnits', 'inches'); set(gcf, 'PaperSize', [5.5 6.5]);

set(gcf, 'PaperPosition', [0 0 5.5 6.5]);
set(gca,'yscale','log','ylim',[370 1200],'ytick',[400, 500, 600, 800, 1200])
saveas(gcf,'Linzetal2021_AoA_Fig6','pdf')
%% Figure 7
% 7a: Overturning comparison
figure
h1=plot(mean(Mutotal2,2),Thlevels(:), 'color',[0.6 0.1 0.8], 'linewidth',2);
hold on
h2 = plot(mean(mutot,2)/365/24/60/60,ideal_levels,'color',[0.03 0.58 0.13], 'linewidth',2);
load('N2O_andrews_overturning.mat')
N2O_averages=[nanmean(totalmass(:,13:12+12)./(extage(:,13:24)-tropage(:,13:24)),2)...
    nanmean(totalmass(:,25:24+12)./(extage(:,25:24+12)-tropage(:,25:24+12)),2)...
    nanmean(totalmass(:,37:36+12)./(extage(:,37:36+12)-tropage(:,37:36+12)),2)...
    nanmean(totalmass(:,49:36+24)./(extage(:,49:36+24)-tropage(:,49:36+24)),2)...
    nanmean(totalmass(:,61:36+36)./(extage(:,61:36+36)-tropage(:,61:36+36)),2)...
    nanmean(totalmass(:,73:36+48)./(extage(:,73:36+48)-tropage(:,73:36+48)),2)...
    nanmean(totalmass(:,85:36+60)./(extage(:,85:36+60)-tropage(:,85:36+60)),2)...
    nanmean(totalmass(:,97:108)./(extage(:,97:108)-tropage(:,97:108)),2)];
N2O_stddev=nanstd(N2O_averages,1,2);
dx1=N2O_stddev(5:22)/365/24/60/60;
x=N2O_overturning(7:10)/365/24/60/60;
N2O_overturning2 = nanmean(totalmass(2:29,13:108)./(extage(2:29,13:108)-tropage(2:29,13:108)),2)/365/24/60/60;
x2 = N2O_overturning2(4:21);
y=Thlevels(5:22);
y1 = Thlevels(7:10);
fill([x2-dx1;flipud(x2+dx1)],[y';flipud(y')],[.8 .8 .8],'linestyle','none');
hold on
h3a=plot(x2,y,'color','k','linewidth',1,'linestyle','-.');
h3=plot(x,y1,'color','k','linewidth',3);
h1=plot(mean(Mutotal2,2),Thlevels(:), 'color',[0.6 0.1 0.8], 'linewidth',2);
hold on
h2 = plot(mean(mutot,2)/365/24/60/60,ideal_levels,'color',[0.03 0.58 0.13], 'linewidth',2);
xlabel('overturning (kg/s)')
ylabel('\Theta (K)')
set(gca,'Xlim',[0 12.5*10^9])
set(gca,'Ylim',[400 1200])
set(gca,'yscale','log')
set(gca,'fontsize',18)
set(gca,'XTick',[0:2.5:12.5]*10^9)
legend([h1 h2 h3],{'WACCM', 'Ideal', 'GOZCARDS N_2O'},'location','northeast','fontsize',22);
legend('boxoff')
hold on
ax1=gca;
set(ax1,'box','off','color','none')
hold on
ax1_pos = get(ax1,'Position'); 
hold on

set(gcf, 'PaperUnits', 'inches'); set(gcf, 'PaperSize', [8 6]);

set(gcf, 'PaperPosition', [0 0 8 6]);
saveas(gcf,'Linzetal2021_AoA_Fig7a','pdf')

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%% Calculate the epsilon for WACCM, for FMS, and for N2O
% epsilon = muin/(dmu/dtheta)
% we just calculated mu(theta)
% now need to calculate muin
% muin= mu/dgamma (dgammaup/dtheta)-sigmau/dgamma

%epsilon_WACCM = muinoversig*diff(Thlevels)/diff(Mutotal2)*sigmaup


dmudtheta_ideal = zeros(size(mutot));
dmudtheta_ideal(1,:) = (mutot(2,:) - mutot(1,:))./(ideal_levels(2)-ideal_levels(1));
for ll = 2:16
    dmudtheta_ideal(ll,:) = (mutot(ll+1) - mutot(ll-1))./(ideal_levels(ll+1)-ideal_levels(ll-1));
end

epsilon_ideal = -muinoversigmau.*sigu./dmudtheta_ideal;
epsilon_ideal_wavgs = -mean(muinoversigmau,2).*mean(sigu,2)./mean(dmudtheta_ideal,2);

% load sigmaup for GOZCARDS_N2O
load('sigmaup_N2O.mat')
muin_N2O = totalmass./(extage - tropage).^2.* dgammaudth - sigmaup_n2o./(extage - tropage);

N2O_overturning3 = nanmean(totalmass(:,13:108)./(extage(:,13:108)-tropage(:,13:108)),2);
dmudtheta_n2o = zeros(size(muin_N2O));
dmudtheta_WACCM = zeros(size(Mutotal2));
dmudtheta_n2o(1,:) = (N2O_overturning3(2,:) - N2O_overturning3(1,:))./(Thlevels(2)-Thlevels(1));
dmudtheta_WACCM(1,:) = (Mutotal2(2,:) - Mutotal2(1,:))./(Thlevels(2)-Thlevels(1));
for ll = 2:29
    dmudtheta_n2o(ll,:) = (N2O_overturning3(ll+1) - N2O_overturning3(ll-1))./(Thlevels(ll+1)-Thlevels(ll-1));
    dmudtheta_WACCM(ll,:) = (Mutotal2(ll+1,:) - Mutotal2(ll-1,:))./(Thlevels(ll+1)-Thlevels(ll-1));
end

muinoversig_noavg=Mutotal2.*dgammaupdtheta./sigmaup./(agedwn2-ageup2)-1./(agedwn2-ageup2);

epsilon_WACCM=-muinoversig.*mean(sigmaup,2)./mean(dmudtheta_WACCM,2)/365/24/60/60;%,Thlevels)
epsilon_n2o=-nanmean(muin_N2O,2)./nanmean(dmudtheta_n2o,2);
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


% 7b: Epsilon comparison
figure
h2=plot(epsilon_ideal_wavgs,ideal_levels,'color',[0.03 0.58 0.13], 'linewidth',2);
hold on
h1=plot(epsilon_WACCM,Thlevels, 'color',[0.6 0.1 0.8], 'linewidth',2);
h4=plot(epsilon_n2o,Thlevels,'color','k','linewidth',1,'linestyle','-.');
y1 = Thlevels(7:10)
h3 = plot(epsilon_n2o(7:10),Thlevels(7:10),'color','k','linewidth',3);
xlabel('\epsilon (unitless)')
ylabel('\Theta (K)')
set(gca,'Xlim',[-0.2 1.3])
set(gca,'Ylim',[380 700])
set(gca,'yscale','log')
set(gca,'fontsize',18)
%set(gca,'XTick',[0:2.5:12.5]*10^9)
legend([h1 h2 h3],{'WACCM', 'Ideal', 'GOZCARDS N_2O'},'location','southwest','fontsize',18);
set(gcf, 'PaperUnits', 'inches'); set(gcf, 'PaperSize', [8 6]);

set(gcf, 'PaperPosition', [0 0 8 6]);
saveas(gcf,'Linzetal2021_AoA_Fig7b','pdf')